use std::collections::BTreeSet;
use std::sync::Arc;

use egui::accesskit::{Role, Toggled};
use egui::{Modifiers, PointerButton};
use egui_kittest::kittest::{NodeT as _, Queryable as _};
use parking_lot::Mutex;
use re_sdk::external::re_log_types::{SetStoreInfo, StoreInfo};
use re_sdk::external::re_tuid::Tuid;
use re_sdk::log::Chunk;
use re_sdk::{
    Component as _, ComponentDescriptor, EntityPath, EntityPathPart, RecordingInfo, StoreId,
    StoreKind,
};
use re_viewer::external::re_chunk::{ChunkBuilder, LatestAtQuery};
use re_viewer::external::re_entity_db::EntityDb;
use re_viewer::external::re_sdk_types;
use re_viewer::external::re_viewer_context::{self, ViewerContext, blueprint_timeline};
use re_viewer::viewer_test_utils::AppTestingExt as _;
use re_viewer::{SystemCommand, SystemCommandSender as _};
use re_viewer_context::ContainerId;
use re_viewport_blueprint::ViewportBlueprint;

// use crate::GetSection;
// use crate::GetSection as _;
use crate::ViewerSection;

// Kittest harness utilities specific to the Rerun app.
pub trait HarnessExt<'h> {
    // Initializes the chuck store with a new, empty recording and blueprint.
    fn init_recording(&mut self);

    // Runs a function with the `ViewerContext` generated by the actual Rerun application.
    // fn run_with_viewer_context(&mut self, func: impl FnOnce(&ViewerContext<'_>) + 'static);
    fn run_with_viewer_context<R: 'static>(
        &mut self,
        func: impl FnOnce(&ViewerContext<'_>) -> R + 'static,
    ) -> R;

    // Removes all views and containers from the current blueprint.
    fn clear_current_blueprint(&mut self);

    // Sets up a new viewport blueprint and saves the new one in the chunk store.
    fn setup_viewport_blueprint<R: 'static>(
        &mut self,
        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint) -> R + 'static,
    ) -> R;

    // Adds a new container to the current blueprint and returns its id.
    fn add_blueprint_container(
        &mut self,
        kind: egui_tiles::ContainerKind,
        parent_container: Option<re_viewer_context::ContainerId>,
    ) -> re_viewer_context::ContainerId;

    // Logs an entity to the active recording.
    fn log_entity(
        &mut self,
        entity_path: impl Into<EntityPath>,
        build_chunk: impl FnOnce(ChunkBuilder) -> ChunkBuilder,
    );

    // Get the position of a node in the UI by its label.
    fn get_panel_position(&mut self, label: &str) -> egui::Rect;

    // Click at a position in the UI.
    fn click_at(&mut self, pos: egui::Pos2);

    // Drag-and-drop functions based on position
    fn drag_at(&mut self, pos: egui::Pos2);
    fn hover_at(&mut self, pos: egui::Pos2);
    fn drop_at(&mut self, pos: egui::Pos2);
    fn right_click_at(&mut self, pos: egui::Pos2);

    // Gets the cursor icon
    fn cursor_icon(&mut self) -> egui::CursorIcon;

    // Changes the value of a dropdown menu.
    fn change_dropdown_value(&mut self, dropdown_label: &str, value: &str);

    // Takes a snapshot of the current app state with good-enough snapshot options.
    fn snapshot_app(&mut self, snapshot_name: &str);

    // Prints the current viewer state.
    fn debug_viewer_state(&mut self);

    // Opens / closes app panels
    fn set_panel_opened(&mut self, panel_label: &str, opened: bool);

    fn set_blueprint_panel_opened(&mut self, opened: bool) {
        self.set_panel_opened("Blueprint panel toggle", opened);
    }

    fn set_selection_panel_opened(&mut self, opened: bool) {
        self.set_panel_opened("Selection panel toggle", opened);
    }

    fn set_time_panel_opened(&mut self, opened: bool) {
        self.set_panel_opened("Time panel toggle", opened);
    }

    // Get a viewer section, eg. blueprint tree or selection panel.
    fn section<'a>(&'a mut self, section_label: &'a str) -> ViewerSection<'a, 'h>
    where
        'h: 'a;

    // The viewer section whose root node is the root of the app.
    fn root_section<'a>(&'a mut self) -> ViewerSection<'a, 'h>
    where
        'h: 'a;

    /// Helper function to save the active recording to file for troubleshooting.
    ///
    /// Note: Right now it _only_ saves the recording and blueprints are ignored.
    fn save_recording_to_file(&mut self, path: impl AsRef<std::path::Path>);

    /// Helper function to save the active blueprint to file for troubleshooting.
    fn save_blueprint_to_file(&mut self, path: impl AsRef<std::path::Path>);

    // The viewer section whose root node is the blueprint tree.
    fn blueprint_tree<'a>(&'a mut self) -> ViewerSection<'a, 'h> {
        self.section("_blueprint_tree")
    }

    // The viewer section whose root node is the streams tree.
    fn streams_tree<'a>(&'a mut self) -> ViewerSection<'a, 'h> {
        self.section("_streams_tree")
    }

    // The viewer section whose root node is the selection panel.
    fn selection_panel<'a>(&'a mut self) -> ViewerSection<'a, 'h> {
        self.section("_selection_panel")
    }

    // Convenience proxy functions to the root section:

    fn click_label(&mut self, label: &str) {
        self.root_section().click_label(label);
    }

    fn right_click_label(&mut self, label: &str) {
        self.root_section().right_click_label(label);
    }

    fn click_label_contains(&mut self, label: &str) {
        self.root_section().click_label_contains(label);
    }

    fn click_nth_label(&mut self, label: &str, index: usize) {
        self.root_section().click_nth_label(label, index);
    }

    fn right_click_nth_label(&mut self, label: &str, index: usize) {
        self.root_section().right_click_nth_label(label, index);
    }

    fn hover_label_contains(&mut self, label: &str) {
        self.root_section().hover_label_contains(label);
    }

    fn hover_nth_label(&mut self, label: &str, index: usize) {
        self.root_section().hover_nth_label(label, index);
    }

    fn drag_nth_label(&mut self, label: &str, index: usize) {
        self.root_section().drag_nth_label(label, index);
    }

    fn drop_nth_label(&mut self, label: &str, index: usize) {
        self.root_section().drop_nth_label(label, index);
    }
}

impl<'h> HarnessExt<'h> for egui_kittest::Harness<'h, re_viewer::App> {
    fn clear_current_blueprint(&mut self) {
        self.setup_viewport_blueprint(|_viewer_context, blueprint| {
            for item in blueprint.contents_iter() {
                blueprint.remove_contents(item);
            }
        });
        self.run_ok();
    }

    fn setup_viewport_blueprint<R: 'static>(
        &mut self,
        setup_blueprint: impl FnOnce(&ViewerContext<'_>, &mut ViewportBlueprint) -> R + 'static,
    ) -> R {
        let result = self.run_with_viewer_context(|viewer_context| {
            let blueprint_query = LatestAtQuery::latest(blueprint_timeline());
            let mut viewport_blueprint =
                ViewportBlueprint::from_db(viewer_context.blueprint_db(), &blueprint_query);
            let result = setup_blueprint(viewer_context, &mut viewport_blueprint);
            viewport_blueprint.save_to_blueprint_store(viewer_context);
            result
        });
        self.run_ok();
        result
    }

    fn run_with_viewer_context<R: 'static>(
        &mut self,
        func: impl FnOnce(&ViewerContext<'_>) -> R + 'static,
    ) -> R {
        let result = Arc::new(Mutex::new(None));
        let result_clone = Arc::clone(&result);
        self.state_mut()
            .testonly_set_test_hook(Box::new(move |viewer_context| {
                *result_clone.lock() = Some(func(viewer_context));
            }));
        self.run_ok();
        result
            .lock()
            .take()
            .expect("test hook should have been called")
    }

    fn log_entity(
        &mut self,
        entity_path: impl Into<EntityPath>,
        build_chunk: impl FnOnce(ChunkBuilder) -> ChunkBuilder,
    ) {
        let app = self.state_mut();
        let builder = build_chunk(Chunk::builder(entity_path));
        let store_hub = app.testonly_get_store_hub();
        let active_recording = store_hub
            .active_recording_mut()
            .expect("active_recording should be initialized");
        active_recording
            .add_chunk(&Arc::new(
                builder.build().expect("chunk should be successfully built"),
            ))
            .expect("chunk should be successfully added");
        self.run_ok();
    }

    fn init_recording(&mut self) {
        let app = self.state_mut();
        let store_hub = app.testonly_get_store_hub();

        let store_info = StoreInfo::testing_with_recording_id("test_recording"); // Fixed id shouldn't cause any problems with store subscribers here since we tear down the entire application for every test.
        let application_id = store_info.application_id().clone();
        let recording_store_id = store_info.store_id.clone();
        let mut recording_store = EntityDb::new(recording_store_id.clone());

        recording_store.set_store_info(SetStoreInfo {
            row_id: Tuid::new(),
            info: store_info,
        });
        {
            // Set RecordingInfo:
            recording_store
                .set_recording_property(
                    EntityPath::properties(),
                    RecordingInfo::descriptor_name(),
                    &re_sdk_types::components::Name::from("Test recording"),
                )
                .expect("Failed to set recording name");
            recording_store
                .set_recording_property(
                    EntityPath::properties(),
                    RecordingInfo::descriptor_start_time(),
                    &re_sdk_types::components::Timestamp::from(0),
                )
                .expect("Failed to set recording start time");
        }
        {
            // Set some custom recording properties:
            recording_store
                .set_recording_property(
                    EntityPath::properties() / EntityPathPart::from("episode"),
                    ComponentDescriptor {
                        archetype: None,
                        component: "location".into(),
                        component_type: Some(re_sdk_types::components::Text::name()),
                    },
                    &re_sdk_types::components::Text::from("Swallow Falls"),
                )
                .expect("Failed to set recording property");
            recording_store
                .set_recording_property(
                    EntityPath::properties() / EntityPathPart::from("episode"),
                    ComponentDescriptor {
                        archetype: None,
                        component: "weather".into(),
                        component_type: Some(re_sdk_types::components::Text::name()),
                    },
                    &re_sdk_types::components::Text::from("Cloudy with meatballs"),
                )
                .expect("Failed to set recording property");
        }

        let blueprint_id = StoreId::random(StoreKind::Blueprint, application_id);
        let blueprint_store = EntityDb::new(blueprint_id.clone());

        store_hub.insert_entity_db(recording_store);
        store_hub.insert_entity_db(blueprint_store);
        store_hub.set_active_recording_id(recording_store_id.clone());
        store_hub
            .set_cloned_blueprint_active_for_app(&blueprint_id)
            .expect("Failed to set blueprint as active");

        app.command_sender.send_system(SystemCommand::set_selection(
            re_viewer_context::Item::StoreId(recording_store_id.clone()),
        ));
        self.run_ok();
    }

    fn get_panel_position(&mut self, label: &str) -> egui::Rect {
        self.get_by_role_and_label(Role::Pane, label).rect()
    }

    fn click_at(&mut self, pos: egui::Pos2) {
        for pressed in [true, false] {
            self.event(egui::Event::PointerButton {
                pos,
                button: PointerButton::Primary,
                pressed,
                modifiers: Modifiers::NONE,
            });
            self.run();
        }
    }

    fn drag_at(&mut self, pos: egui::Pos2) {
        self.event(egui::Event::PointerButton {
            pos,
            button: PointerButton::Primary,
            pressed: true,
            modifiers: Modifiers::NONE,
        });
        self.run_ok();
    }

    fn hover_at(&mut self, pos: egui::Pos2) {
        self.event(egui::Event::PointerMoved(pos));
        self.run_ok();
    }

    fn drop_at(&mut self, pos: egui::Pos2) {
        self.event(egui::Event::PointerButton {
            pos,
            button: PointerButton::Primary,
            pressed: false,
            modifiers: Modifiers::NONE,
        });
        self.remove_cursor();
        self.run_ok();
    }

    fn right_click_at(&mut self, pos: egui::Pos2) {
        self.event(egui::Event::PointerButton {
            pos,
            button: PointerButton::Secondary,
            pressed: true,
            modifiers: Modifiers::NONE,
        });
        self.event(egui::Event::PointerButton {
            pos,
            button: PointerButton::Secondary,
            pressed: false,
            modifiers: Modifiers::NONE,
        });
        self.run();
    }

    fn cursor_icon(&mut self) -> egui::CursorIcon {
        self.run_with_viewer_context(|viewer_context| {
            viewer_context.egui_ctx().output(|o| o.cursor_icon)
        })
    }

    fn debug_viewer_state(&mut self) {
        println!(
            "Active recording: {:#?}",
            self.state_mut().testonly_get_store_hub().active_recording()
        );
        println!(
            "Active blueprint: {:#?}",
            self.state_mut().testonly_get_store_hub().active_blueprint()
        );
        self.setup_viewport_blueprint(|_viewer_context, blueprint| {
            println!("Blueprint view count: {}", blueprint.views.len());
            for id in blueprint.view_ids() {
                println!("View id: {id}");
            }
            println!(
                "Display mode: {:?}",
                _viewer_context.global_context.display_mode
            );
        });
    }

    fn change_dropdown_value(&mut self, dropdown_label: &str, value: &str) {
        let node = self.get_by_role_and_label(Role::ComboBox, dropdown_label);
        node.click();
        self.run_ok();
        self.click_label(value);
    }

    fn snapshot_app(&mut self, snapshot_name: &str) {
        self.run_ok();
        self.snapshot(snapshot_name);
    }

    fn add_blueprint_container(
        &mut self,
        kind: egui_tiles::ContainerKind,
        parent_container: Option<re_viewer_context::ContainerId>,
    ) -> ContainerId {
        // Collect list of container ids before adding new container
        let old_container_ids: BTreeSet<_> =
            self.setup_viewport_blueprint(|_viewer_context, blueprint| {
                blueprint.containers.keys().copied().collect()
            });

        // Add new container
        self.setup_viewport_blueprint(move |_viewer_context, blueprint| {
            blueprint.add_container(kind, parent_container);
        });

        // Collect list of container ids after adding new container
        let new_container_ids: BTreeSet<_> =
            self.setup_viewport_blueprint(|_viewer_context, blueprint| {
                blueprint.containers.keys().copied().collect()
            });
        self.run_ok();

        // Find the new container id
        let container_ids = new_container_ids
            .difference(&old_container_ids)
            .collect::<Vec<_>>();
        assert!(
            container_ids.len() == 1,
            "Expected one new container id, got {container_ids:?}"
        );
        *container_ids[0]
    }

    fn set_panel_opened(&mut self, panel_label: &str, opened: bool) {
        let node = self.get_by_label(panel_label);
        let is_open = Some(Toggled::True) == node.accesskit_node().data().toggled();
        if is_open != opened {
            self.click_label(panel_label);
        }
        self.remove_cursor();
        self.run_ok();
    }

    fn section<'a>(&'a mut self, section_label: &'a str) -> ViewerSection<'a, 'h>
    where
        'h: 'a,
    {
        ViewerSection::<'a, 'h> {
            harness: self,
            section_label: Some(section_label),
        }
    }

    fn root_section<'a>(&'a mut self) -> ViewerSection<'a, 'h>
    where
        'h: 'a,
    {
        ViewerSection::<'a, 'h> {
            harness: self,
            section_label: None,
        }
    }

    fn save_recording_to_file(&mut self, path: impl AsRef<std::path::Path>) {
        let mut file = std::fs::File::create(&path)
            .unwrap_or_else(|e| panic!("Failed to create file at {:?}: {}", path.as_ref(), e));

        let store_hub = self.state_mut().testonly_get_store_hub();
        let recording_entity_db = store_hub.active_recording().expect("No active recording");
        let messages = recording_entity_db.to_messages(None);

        let encoding_options = re_log_encoding::rrd::EncodingOptions::PROTOBUF_COMPRESSED;
        re_log_encoding::Encoder::encode_into(
            re_build_info::CrateVersion::LOCAL,
            encoding_options,
            messages,
            &mut file,
        )
        .expect("Failed to encode recording to file");
    }

    fn save_blueprint_to_file(&mut self, path: impl AsRef<std::path::Path>) {
        let mut file = std::fs::File::create(&path)
            .unwrap_or_else(|e| panic!("Failed to create file at {:?}: {}", path.as_ref(), e));

        let store_hub = self.state_mut().testonly_get_store_hub();
        let blueprint_entity_db = store_hub.active_blueprint().expect("No active blueprint");
        let messages = blueprint_entity_db.to_messages(None);

        let encoding_options = re_log_encoding::rrd::EncodingOptions::PROTOBUF_COMPRESSED;
        re_log_encoding::Encoder::encode_into(
            re_build_info::CrateVersion::LOCAL,
            encoding_options,
            messages,
            &mut file,
        )
        .expect("Failed to encode blueprint to file");
    }
}
